package org.springframework.session.web.socket.server;

import java.time.Instant;
import java.util.EnumSet;
import java.util.Map;
import java.util.Set;

import javax.servlet.http.HttpSession;

import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
import org.springframework.messaging.simp.SimpMessageType;
import org.springframework.messaging.support.ChannelInterceptor;
import org.springframework.session.Session;
import org.springframework.session.SessionRepository;
import org.springframework.util.Assert;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

public final class SessionRepositoryMessageInterceptor<S extends Session> implements ChannelInterceptor, HandshakeInterceptor {

	private static final String SPRING_SESSION_ID_ATTR_NAME = "SPRING.SESSION.ID";

	private final SessionRepository<S> sessionRepository;

	private Set<SimpMessageType> matchingMessageTypes;

	public SessionRepositoryMessageInterceptor(SessionRepository<S> sessionRepository) {
		Assert.notNull(sessionRepository, "sessionRepository cannot be null");
		this.sessionRepository = sessionRepository;
		this.matchingMessageTypes = EnumSet.of(SimpMessageType.CONNECT, SimpMessageType.MESSAGE, SimpMessageType.SUBSCRIBE, SimpMessageType.UNSUBSCRIBE);
	}

	public void setMatchingMessageTypes(Set<SimpMessageType> matchingMessageTypes) {
		Assert.notEmpty(matchingMessageTypes, "matchingMessageTypes cannot be null or empty");
		this.matchingMessageTypes = matchingMessageTypes;
	}

	@Override
	public Message<?> preSend(Message<?> message, MessageChannel channel) {
		if (message == null) {
			return message;
		}
		SimpMessageType messageType = SimpMessageHeaderAccessor.getMessageType(message.getHeaders());
		if (!this.matchingMessageTypes.contains(messageType)) {
			return message;
		}
		Map<String, Object> sessionHeaders = SimpMessageHeaderAccessor.getSessionAttributes(message.getHeaders());
		String sessionId = (sessionHeaders != null) ? (String) sessionHeaders.get(SPRING_SESSION_ID_ATTR_NAME) : null;
		if (sessionId != null) {
			S session = this.sessionRepository.findById(sessionId);
			if (session != null) {
				// update the last accessed time
				session.setLastAccessedTime(Instant.now());
				this.sessionRepository.save(session);
			}
		}
		return message;
	}

	@Override
	public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) {
		if (request instanceof ServletServerHttpRequest) {
			ServletServerHttpRequest servletRequest = (ServletServerHttpRequest) request;
			HttpSession session = servletRequest.getServletRequest().getSession(false);
			if (session != null) {
				setSessionId(attributes, session.getId());
			}
		}
		return true;
	}

	@Override
	public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception exception) {
	}

	public static String getSessionId(Map<String, Object> attributes) {
		return (String) attributes.get(SPRING_SESSION_ID_ATTR_NAME);
	}

	public static void setSessionId(Map<String, Object> attributes, String sessionId) {
		attributes.put(SPRING_SESSION_ID_ATTR_NAME, sessionId);
	}
}
